package com.imut.lagain.websocket;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.imut.lagain.util.JwtUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;

import java.io.IOException;
import java.net.URI;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * WebSocket处理器
 */
@Component
public class WebSocketHandler implements org.springframework.web.socket.WebSocketHandler {
    private static final Logger log = LoggerFactory.getLogger(WebSocketHandler.class);
    
    @Autowired
    private JwtUtil jwtUtil;
    
    private final ObjectMapper objectMapper = new ObjectMapper();
    
    // 存储用户ID与WebSocket会话的映射
    private final Map<Long, WebSocketSession> userSessions = new ConcurrentHashMap<>();
    
    // 存储会话ID与用户ID的映射
    private final Map<String, Long> sessionUsers = new ConcurrentHashMap<>();
    
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        try {
            // 从查询参数中获取token
            URI uri = session.getUri();
            String query = uri.getQuery();
            String token = null;
            
            if (query != null) {
                String[] params = query.split("&");
                for (String param : params) {
                    String[] keyValue = param.split("=");
                    if (keyValue.length == 2 && "token".equals(keyValue[0])) {
                        token = keyValue[1];
                        break;
                    }
                }
            }
            
            if (token == null || token.isEmpty()) {
                log.warn("WebSocket连接缺少token参数");
                session.close(CloseStatus.BAD_DATA.withReason("Missing token"));
                return;
            }
            
            // 验证token并获取用户ID
            Long userId = jwtUtil.getUserIdFromToken(token);
            if (userId == null) {
                log.warn("WebSocket连接token无效");
                session.close(CloseStatus.BAD_DATA.withReason("Invalid token"));
                return;
            }
            
            // 存储用户会话映射
            userSessions.put(userId, session);
            sessionUsers.put(session.getId(), userId);
            
            log.info("用户{}建立WebSocket连接，会话ID: {}", userId, session.getId());
            
            // 发送连接成功消息
            sendMessage(session, createMessage("connection", "连接成功", null));
            
        } catch (Exception e) {
            log.error("WebSocket连接建立失败", e);
            session.close(CloseStatus.SERVER_ERROR.withReason("Connection failed"));
        }
    }
    
    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        try {
            String payload = message.getPayload().toString();
            log.debug("收到WebSocket消息: {}", payload);
            
            // 解析消息
            Map<String, Object> messageData = objectMapper.readValue(payload, Map.class);
            String type = (String) messageData.get("type");
            
            Long userId = sessionUsers.get(session.getId());
            if (userId == null) {
                log.warn("未找到会话对应的用户ID: {}", session.getId());
                return;
            }
            
            // 处理不同类型的消息
            switch (type) {
                case "ping":
                    // 心跳检测
                    sendMessage(session, createMessage("pong", "心跳响应", null));
                    break;
                case "subscribe":
                    // 订阅特定类型的通知
                    handleSubscribe(session, userId, messageData);
                    break;
                default:
                    log.warn("未知的消息类型: {}", type);
                    break;
            }
            
        } catch (Exception e) {
            log.error("处理WebSocket消息失败", e);
        }
    }
    
    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        log.error("WebSocket传输错误，会话ID: {}", session.getId(), exception);
        cleanupSession(session);
    }
    
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        log.info("WebSocket连接关闭，会话ID: {}，状态: {}", session.getId(), closeStatus);
        cleanupSession(session);
    }
    
    @Override
    public boolean supportsPartialMessages() {
        return false;
    }
    
    /**
     * 向指定用户发送消息
     */
    public void sendMessageToUser(Long userId, String type, String content, Object data) {
        WebSocketSession session = userSessions.get(userId);
        if (session != null && session.isOpen()) {
            try {
                String message = createMessage(type, content, data);
                sendMessage(session, message);
            } catch (Exception e) {
                log.error("向用户{}发送WebSocket消息失败", userId, e);
            }
        } else {
            log.debug("用户{}的WebSocket会话不存在或已关闭", userId);
        }
    }
    
    /**
     * 向指定用户发送新消息通知
     */
    public void sendNewMessageNotification(Long userId, Map<String, Object> messageData) {
        sendMessageToUser(userId, "new_message", "收到新消息", messageData);
    }
    
    /**
     * 向指定用户发送系统通知
     */
    public void sendSystemNotification(Long userId, Map<String, Object> notificationData) {
        sendMessageToUser(userId, "system_notification", "收到系统通知", notificationData);
    }
    
    /**
     * 向指定用户发送生理期提醒
     */
    public void sendPeriodReminder(Long userId, Map<String, Object> reminderData) {
        sendMessageToUser(userId, "period_reminder", "生理期提醒", reminderData);
    }
    
    /**
     * 获取在线用户数量
     */
    public int getOnlineUserCount() {
        return userSessions.size();
    }
    
    /**
     * 检查用户是否在线
     */
    public boolean isUserOnline(Long userId) {
        WebSocketSession session = userSessions.get(userId);
        return session != null && session.isOpen();
    }
    
    /**
     * 处理订阅请求
     */
    private void handleSubscribe(WebSocketSession session, Long userId, Map<String, Object> messageData) {
        try {
            String subscribeType = (String) messageData.get("subscribeType");
            log.info("用户{}订阅通知类型: {}", userId, subscribeType);
            
            // 这里可以根据需要实现特定的订阅逻辑
            sendMessage(session, createMessage("subscribe_success", "订阅成功", 
                Map.of("subscribeType", subscribeType)));
            
        } catch (Exception e) {
            log.error("处理订阅请求失败", e);
        }
    }
    
    /**
     * 清理会话
     */
    private void cleanupSession(WebSocketSession session) {
        Long userId = sessionUsers.remove(session.getId());
        if (userId != null) {
            userSessions.remove(userId);
            log.info("清理用户{}的WebSocket会话", userId);
        }
    }
    
    /**
     * 发送消息到会话
     */
    private void sendMessage(WebSocketSession session, String message) throws IOException {
        if (session.isOpen()) {
            session.sendMessage(new TextMessage(message));
        }
    }
    
    /**
     * 创建消息JSON字符串
     */
    private String createMessage(String type, String content, Object data) {
        try {
            Map<String, Object> message = Map.of(
                "type", type,
                "content", content,
                "data", data != null ? data : Map.of(),
                "timestamp", System.currentTimeMillis()
            );
            return objectMapper.writeValueAsString(message);
        } catch (Exception e) {
            log.error("创建消息JSON失败", e);
            return "{\"type\":\"error\",\"content\":\"消息格式错误\"}";
        }
    }
}